"""
数据集类
"""
import sklearn
from sklearn import datasets
import numpy as np
import matplotlib.pyplot as plt

from dezero.transforms import Compose, Flatten, ToFloat, Normalize

class Dataset:
    def __init__(self, train=True, transform=None, target_transform=None):
        """
        数据集基类
        参数：
            train：(bool) 区别是训练还是测试
            transform: (func) 对单个输入数据进行转换处理 默认直接输出
            target_transform: (func) 对单个标签进行转换处理
        """
        self.train = train
        self.transform = transform
        self.target_transform = target_transform
        if self.transform is None:
            self.transform = lambda x:x
            #lambda x:x 相当于不进行转换直接输出
        if self.target_transform is None:
            self.target_transform = lambda x:x

        self.data = None    #保存输入数据
        self.label = None   #保存标签
        self.prepare()      #准备数据，在具体类实现

    def __getitem__(self, index):
        assert np.isscalar(index)   #只支持index是标量情况,保护支持切片操作
        if self.label is None:
            return self.transform(self.data[index]),None
        else:
            return self.transform(self.data[index]),\
                self.target_transform(self.label[index])

    def __len__(self):
        return len(self.data)
    def prepare(self):
        """数据预处理由具体类实现"""
        pass

#======================具体数据集=============
#螺旋数据集
def get_spiral(train=True):
    """准备螺旋数据集
    参数:
        train: 是否返回训练数据集，否则返回测试数据
    返回:
        x:输入数据
        t:训练数据(标签) 3类分类问题 所以元素值为0、1、2
    """
    seed = 1984 if train else 2020
    np.random.seed(seed=seed)

    num_data, num_class, input_dim = 100, 3, 2
    data_size = num_class * num_data
    x = np.zeros((data_size, input_dim), dtype=np.float32)
    t = np.zeros(data_size, dtype=np.int32)

    for j in range(num_class):
        for i in range(num_data):
            rate = i / num_data
            radius = 1.0 * rate
            theta = j * 4.0 + 4.0 * rate + np.random.randn() * 0.2
            ix = num_data * j + i
            x[ix] = np.array([radius * np.sin(theta),
                              radius * np.cos(theta)]).flatten()
            t[ix] = j
    # Shuffle
    indices = np.random.permutation(num_data * num_class)
    x = x[indices]
    t = t[indices]
    return x, t

class Spiral(Dataset):
    def prepare(self):
        self.data, self.label = get_spiral(self.train)


class BigData(Dataset):
    def __getitem__(self,index):
        x = np.load('data/{}.npy'.format(index))
        t = np.load('label/{}.npy'.format(index))
        return x, t
    def __len__(self):
        return 1000000

#==============================
# MNIST 数据集
#==============================
import gzip
import os

class MNIST(Dataset):
    """从文件获取"""
    def __init__(self, train=True,
                 transform=Compose([Flatten(), ToFloat(),
                                     Normalize(0., 255.)]),
                 target_transform=None):
        super().__init__(train, transform, target_transform)

    def prepare(self):
        url = os.path.dirname(__file__)+'\\data\\'
        train_files = {'target': 'train-images-idx3-ubyte.gz',
                       'label': 'train-labels-idx1-ubyte.gz'}
        test_files = {'target': 't10k-images-idx3-ubyte.gz',
                      'label': 't10k-labels-idx1-ubyte.gz'}

        files = train_files if self.train else test_files
        data_path = url + files['target']
        label_path = url + files['label']

        self.data = self._load_data(data_path)
        self.label = self._load_label(label_path)

    def _load_label(self, filepath):
        with gzip.open(filepath, 'rb') as f:
            labels = np.frombuffer(f.read(), np.uint8, offset=8)
        return labels

    def _load_data(self, filepath):
        with gzip.open(filepath, 'rb') as f:
            data = np.frombuffer(f.read(), np.uint8, offset=16)
        data = data.reshape(-1, 1, 28, 28)
        return data

    def show(self, row=10, col=10):
        H, W = 28, 28
        img = np.zeros((H * row, W * col))
        for r in range(row):
            for c in range(col):
                img[r * H:(r + 1) * H, c * W:(c + 1) * W] = self.data[ r+c ].reshape(H, W)
        plt.imshow(img, cmap='gray', interpolation='nearest')
        plt.axis('off')
        plt.show()

    @staticmethod
    def labels():
        return {0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9'}

class MNIST_sklearn(Dataset):
    """从sklearn获取"""
    def __init__(self, train=True,
                 transform=Compose([Flatten(), ToFloat(),
                                     Normalize(0., 255.)]),
                 target_transform=None):
        super().__init__(train, transform, target_transform)

    def prepare(self):
        digits = sklearn.datasets.load_digits()

        self.data = digits['data'].reshape(-1, 1, 8, 8) #shape(1796,64)
        self.label = digits['target']

    def show(self, row=10, col=10):
        H, W = 8, 8
        img = np.zeros((H * row, W * col))
        for r in range(row):
            for c in range(col):
                img[r * H:(r + 1) * H, c * W:(c + 1) * W] = self.data[ r+c ].reshape(H, W)
        plt.imshow(img, cmap='gray', interpolation='nearest')
        plt.axis('off')
        plt.show()

    @staticmethod
    def labels():
        return {0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9'}
